import numpy as np
import matplotlib.pyplot as plt


class TMM:
    """
    Initialize a transfer matrix system.
    Parameters:
        input_t_layers: array, thickness of each layer
        input_n_layers: array, refractive index of each layer
        unit: string, unit of frequency and length, in either 'nm/THz' or 'mm/GHz'

    """

    def __init__(self, input_t_layers, input_n_layers, unit='nm/THz', n_top=1, n_bottom=1):
        if len(input_t_layers) != len(input_n_layers):
            raise ValueError('The number of layers and refractive indices do not match.')

        self.t_layers = input_t_layers
        self.n_layers = input_n_layers

        self.unit = unit
        self.n_top = n_top
        self.n_bottom = n_bottom
        self.num_layers = len(self.t_layers)

        if self.unit == 'nm/THz':
            self.c = 2.998e5
        else:
            self.c = 2.998e2

    def calculate_tm(self, input_freq):
        """
        Calculate the transfer matrix system.
        Parameters:
            input_freq: float, frequency where the TMM is calculated, in the chosen unit
        Returns:
            Q: transfer matrix of the system under freq

        """
        z_layers = 1 / self.n_layers
        z_top = 1 / self.n_top
        z_bottom = 1 / self.n_bottom
        q = np.eye(2, dtype=float)
        for i in range(self.num_layers):
            if i == 0:
                z_rel = z_layers[0] / z_top
                q = np.dot(np.array([[1 + z_rel, 1 - z_rel], [1 - z_rel, 1 + z_rel]]), q) / 2
            else:
                z_rel = z_layers[i] / z_layers[i - 1]
                q = np.dot(np.array([[1 + z_rel, 1 - z_rel], [1 - z_rel, 1 + z_rel]]), q) / 2

            k = 2 * np.pi * input_freq / self.c
            phase = self.n_layers[i] * k * self.t_layers[i]
            q = np.dot(np.diag([np.exp(1j * phase), np.exp(-1j * phase)]), q)

        z_rel = z_bottom / z_layers[-1]
        q = np.dot(np.array([[1 + z_rel, 1 - z_rel], [1 - z_rel, 1 + z_rel]]), q) / 2

        return q

    def calculate_spectrum(self, freq_start, freq_end, num_freq):
        """
        Calculate the transmission and reflection of the system.
        Parameters:
            freq_start: float, start frequency of the spectrum, in the chosen unit
            freq_end: float, end frequency of the spectrum, in the chosen unit
            num_freq: int, number of frequency points in the spectrum
        Returns:
            frequency: array, frequency points of the spectrum
            t: array, transmission of the spectrum
            r: array, reflection of the spectrum
        """
        out_freq = np.linspace(freq_start, freq_end, num_freq)
        tran = np.zeros(num_freq, dtype=complex)
        refl = np.zeros(num_freq, dtype=complex)
        for i in range(num_freq):
            q = self.calculate_tm(out_freq[i])
            tran[i] = q[0, 0] - q[0, 1] * q[1, 0] / q[1, 1]
            refl[i] = -q[1, 0] / q[1, 1]

        return out_freq, tran, refl


if __name__ == '__main__':
    t_layers = np.array([10] * 20)
    n_layers = np.array([1.45, 2] * (20 // 2))
    tmm = TMM(t_layers, n_layers, unit='nm/THz')
    _, t, _ = tmm.calculate_spectrum(450, 750, 200)
    plt.plot(np.abs(t) ** 2)
    plt.show()
